from jax import numpy as jnp
from jax import random
from jax import jit
from python_basic.study.python_cookbook.ch9_metaprog.x91_timethis import timethis
from PyCmpltrtok.common import sep, MyTimer


@timethis
def selu(x, alpha=1.67, lmbda=1.05):
    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)


sep('Generaste large vector')
key = random.PRNGKey(0)
x = random.normal(key, (1000000,))
xN = 10

sep('selu xN')
for _ in range(xN):
    selu(x).block_until_ready()

sep('jit(selu) xN')
jit_selu = jit(selu)
for _ in range(xN):
    with MyTimer() as timer:
        jit_selu(x).block_until_ready()
